import matplotlib.patches as pat
import matplotlib.pyplot as plt
import torch
import numpy as np

from torchvision.transforms import transforms
from darknet import Darknet
from dataset import CocoDataset

import os
from PIL import Image, ImageOps

# transform = transforms.Compose([
#     transforms.Resize(size=(416,416)),
#     transforms.ToTensor(),
# ])
#
# # load模型
# model = Darknet(in_chanel=3,use_fc=False)
# model.load_state_dict(torch.load("./ep10_yolo.ckpt"))
# model.eval()
# #print(model)
#
# dataset =  CocoDataset("./sub_set/train2014","./sub_set/labels/train2014", to_square=True, transform=transform)
# dataloader = torch.utils.data.DataLoader(
#     dataset,
#     batch_size=1,
#     shuffle=True,
#     pin_memory=True,
#     collate_fn=dataset.collate_fn,
# )
#
# for idx, batch in enumerate(dataloader):
#     img_tensor, label = batch
#     break
#
# batch_idx = np.random.randint(0,img_tensor.shape[0]-1) if img_tensor.shape[0]>1 else 0
# img = (img_tensor*255).byte()[batch_idx,...].permute(1,2,0).contiguous()
#
# # 预测
# y_hat,_ = model(img_tensor,label)
# print(_)

transform = transforms.Compose([
    transforms.Resize(size=(416,416)),
    transforms.ToTensor(),
])

# 随机读取一张图片
img = Image.open("./a.jfif")

# pad成正方形
w, h = img.width, img.height
pad_tuple = ( 0, (w-h)//2, 0, (w-h)//2) \
    if w>h else ( abs(w-h)//2, 0, abs(w-h)//2, 0)
img = ImageOps.expand(img, pad_tuple, fill="black")

# 缩放到416*416
img_tensor = transform(img).unsqueeze(0)



# load模型
model = Darknet(in_chanel=3, use_fc=False)
model.load_state_dict(torch.load("./ep10_yolo.ckpt"))
model.eval() # 进入预测模式

# 预测
y_hat,_ = model(img_tensor,None)
print(_)

# nms处理

conf_tresh = 0.8
nms_tresh = 0.4
y = y_hat[0,...]
y = y[y[:,4]>conf_tresh,:] # 取置信度较大的部分

y = y[torch.argsort(y[:,4])] # 根据置信度排序
drop_set = []
reserve_set = [i for i in range(y.shape[0])]
# nms处理
for idx, i in enumerate(reserve_set):
    b1_x1, b1_y1 = y[i][0] - y[i][2] / 2, y[i][1] - y[i][3]
    b1_x2, b1_y2 = y[i][0] + y[i][2] / 2, y[i][1] + y[i][3]

    for j in reserve_set[idx+1:]:
        b2_x1, b2_y1 = y[j][0] - y[j][2] / 2, y[j][1] - y[j][3]
        b2_x2, b2_y2 = y[j][0] + y[j][2] / 2, y[j][1] + y[j][3]

        inter_x1 = max(b1_x1, b2_x1)
        inter_y1 = max(b1_y1, b2_y1)
        inter_x2 = min(b1_x2, b1_x2)
        inter_y2 = min(b1_y2, b2_y2)
        # 计算交集
        I = max(inter_x2 - inter_x1, 0) * max(inter_y2 - inter_y1, 0)
        # 计算并集
        U = (b1_y2 - b1_y1 + 1) * (b1_x2 - b1_x1 + 1) + (b2_y2 - b2_y1 + 1) * (b2_x2 - b2_x1 + 1) -I

        if I/U > nms_tresh:
            drop_set.append(j)
            reserve_set.remove(j)

# 展示结果
img = (img_tensor*255).byte()[0,...].permute(1,2,0).contiguous()
img = np.array(img) #
img = np.array(img) #
plt.imshow(img)
plt.show()
coco_name = open("./coco.names", "r").readlines()

plt.figure()
fig, ax = plt.subplots(1)
ax.imshow(img)

for idx, i in enumerate(range(y.shape[0])):
    if i not in reserve_set:
       continue
    x1, y1 = y[i][0] - y[i][2]/2, y[i][1] - y[i][3]/2
    x2, y2 = y[i][0] + y[i][2]/2, y[i][1] + y[i][3]/2
    x1, y1 = np.clip(int(x1.detach().numpy()),0,416) , np.clip(int(y1.detach().numpy()),0,416)
    x2, y2 = np.clip(int(x2.detach().numpy()),0,416), np.clip(int(y2.detach().numpy()),0,416)

    # !! 5：-1注意改掉
    if torch.max(y[i][5:],dim=0).values < 0.5:
        color = "black"
    else:
        color = "white"
    cls_pred = coco_name[torch.argmax(y[i][5:], dim=0)]
    cmap = plt.get_cmap("tab20b")
    colors = [cmap(i) for i in np.linspace(0, 1, 20)]
    bbox = pat.Rectangle((x1, y1), x2-x1, y2-y1, linewidth=2, edgecolor=colors[idx*3%len(colors)], facecolor="none")
    ax.add_patch(bbox)
    plt.text(
        x1,
        y1-10,
        s=cls_pred+str(torch.max(y[i][5:-1],dim=0).values.detach().numpy()),
            #",yolo"+str(y[i][-1].detach().cpu().numpy()),
        color=color,
        verticalalignment="baseline",
        bbox={"color": colors[idx*3%len(colors)], "pad": 0},
        fontdict={'size': 8, 'color': 'red'}
    )

plt.show()

# plt.figure()
# fig, ax = plt.subplots(1)
# ax.imshow(img)
# img_arr = img.copy()
#
# label = label[:,1:]
# for i in range(label.shape[0]):
#     x1, y1 = int(label[i,1]*416 - label[i,3]*416/2), int(label[i,2]*416 - label[i,4]*416/2)
#     x2, y2 = int(label[i,1]*416 + label[i,3]*416/2), int(label[i,2]*416 + label[i,4]*416/2)
#     bbox = pat.Rectangle((x1, y1), x2-x1, y2-y1, linewidth=2, edgecolor=colors[idx*3%len(colors)], facecolor="none")
#     ax.add_patch(bbox)
#     cls = coco_name[int(label[i,0])]
#     plt.text(
#         x1,
#         y1-10,
#         s=cls,
#         color="white",
#         verticalalignment="baseline",
#         bbox={"color": colors[idx*3%len(colors)], "pad": 0},
#         fontdict={'size': 8, 'color': 'red'}
#     )
#     #cv2.rectangle(img_arr,(x1,y1),(x2,y2),(255,0,0),2)
#
# plt.imshow(img_arr)
# plt.show()





